import time
from rest_framework.test import APIClient
from rest_framework_simplejwt.tokens import RefreshToken

class JWTTestClient:
    def __init__(self):
        self.client = APIClient()
        self.current_access = None
        self.current_refresh = None
    
    def login(self, user):
        refresh = RefreshToken.for_user(user)
        self.current_access = str(refresh.access_token)
        self.current_refresh = str(refresh)
        self.client.credentials(
            HTTP_AUTHORIZATION=f'Bearer {self.current_access}')
        
    def wait_for_expiry(self, token_type='access'):
        """等待token过期"""
        if token_type == 'access':
            time.sleep(6)  # 超过5秒
        else:
            time.sleep(11)  # 超过10秒
    
    def auto_refresh(self):
        """自动化刷新流程"""
        from rest_framework import status
        response = self.client.post(
            '/api/token/refresh/',
            {'refresh': self.current_refresh},
            format='json'
        )
        if response.status_code == status.HTTP_200_OK:
            self.current_access = response.data['access']
            self.client.credentials(
                HTTP_AUTHORIZATION=f'Bearer {self.current_access}')
        return response
